#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>

// JPEG markers
#define JPEG_SOI 0xFFD8
#define JPEG_EOI 0xFFD9
#define JPEG_APP1 0xFFE1
#define JPEG_SOS 0xFFDA

// Exif tags
#define TAG_MAKE  0x010F
#define TAG_MODEL 0x0110
#define TAG_EXIF_IFD_POINTER 0x8769
#define TAG_GPS_IFD_POINTER  0x8825

// ExifIFD tags
#define TAG_ISO_SPEED_RATINGS 0x8827

// GPS tags
#define TAG_GPS_VERSION_ID       0x0000
#define TAG_GPS_LATITUDE_REF     0x0001
#define TAG_GPS_LATITUDE         0x0002
#define TAG_GPS_LONGITUDE_REF    0x0003
#define TAG_GPS_LONGITUDE        0x0004

// TIFF types
#define TYPE_BYTE   1
#define TYPE_ASCII  2
#define TYPE_SHORT  3
#define TYPE_LONG   4
#define TYPE_RATIONAL 5

// Helpers to write big-endian to file
static void write_be16(FILE *f, uint16_t v) {
    uint8_t b[2] = { (uint8_t)((v >> 8) & 0xFF), (uint8_t)(v & 0xFF) };
    fwrite(b, 1, 2, f);
}

// Helpers to append little-endian into memory buffer
static void append_le16(uint8_t *buf, size_t *idx, uint16_t v) {
    buf[(*idx)++] = (uint8_t)(v & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 8) & 0xFF);
}
static void append_le32(uint8_t *buf, size_t *idx, uint32_t v) {
    buf[(*idx)++] = (uint8_t)(v & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 8) & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 16) & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 24) & 0xFF);
}

static void append_bytes(uint8_t *buf, size_t *idx, const void *src, size_t len) {
    memcpy(buf + *idx, src, len);
    *idx += len;
}

// Build Exif payload (content after APP1 length): "Exif\0\0" + TIFF
// Returns payload size in bytes
static size_t build_exif_payload(uint8_t *payload,
                                const char *make,
                                const char *model,
                                double latitude,
                                double longitude,
                                uint16_t iso) {
    size_t idx = 0;
    // "Exif\0\0"
    const uint8_t exif_id[6] = { 'E','x','i','f',0,0 };
    append_bytes(payload, &idx, exif_id, 6);

    // TIFF header (little-endian)
    // Byte order: "II"
    append_bytes(payload, &idx, (const uint8_t*)"II", 2);
    // TIFF version 0x002A
    append_le16(payload, &idx, 0x002A);
    // Offset to 0th IFD from TIFF start (after byte order), which is 8 bytes
    append_le32(payload, &idx, 8);

    // From now on, offsets are relative to TIFF start (position of the 'I' in "II").
    // TIFF start offset within payload is 6 bytes from payload base.

    // Layout planning
    const uint32_t tiff_start = 6; // index in payload where TIFF starts
    const uint32_t ifd0_offset = 8; // relative to TIFF start

    const size_t make_len = strlen(make) + 1;
    const size_t model_len = strlen(model) + 1;

    // IFD0 has 4 entries: Make, Model, ExifIFDPointer, GPSInfoIFDPointer
    const uint32_t ifd0_size = 2 + 4 * 12 + 4; // count + entries + nextIFD
    uint32_t data_offset = ifd0_offset + ifd0_size; // where variable data starts

    // Place Make and Model strings
    const uint32_t make_offset_rel = data_offset; // relative to TIFF start
    data_offset += (uint32_t)make_len;
    const uint32_t model_offset_rel = data_offset;
    data_offset += (uint32_t)model_len;

    // Plan ExifIFD
    const uint32_t exif_ifd_offset_rel = data_offset;
    const uint32_t exif_ifd_size = 2 + 1 * 12 + 4;
    data_offset += exif_ifd_size;

    // Plan GPS IFD (5 entries: VersionID, LatRef, Lat, LonRef, Lon)
    const uint32_t gps_ifd_offset_rel = data_offset;
    const uint32_t gps_ifd_size = 2 + 5 * 12 + 4;
    data_offset += gps_ifd_size;

    // Plan GPS data blocks: latitude and longitude RATIONAL[3] (each 24 bytes)
    const uint32_t gps_lat_data_offset_rel = data_offset;
    data_offset += 24;
    const uint32_t gps_lon_data_offset_rel = data_offset;
    data_offset += 24;

    // Now write IFD0 at payload position tiff_start + ifd0_offset
    // Ensure idx currently equals tiff_start + ifd0_offset
    // We have already appended 6 (Exif ID) + 8 (TIFF header) = 14 bytes
    // We need to append padding if idx isn't at expected position
    size_t expected_ifd0_start = tiff_start + ifd0_offset;
    if (idx < expected_ifd0_start) {
        size_t pad = expected_ifd0_start - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }

    // IFD0 entry count
    append_le16(payload, &idx, 4);

    // Entry: Make (ASCII)
    append_le16(payload, &idx, TAG_MAKE);
    append_le16(payload, &idx, TYPE_ASCII);
    append_le32(payload, &idx, (uint32_t)make_len);
    append_le32(payload, &idx, make_offset_rel);

    // Entry: Model (ASCII)
    append_le16(payload, &idx, TAG_MODEL);
    append_le16(payload, &idx, TYPE_ASCII);
    append_le32(payload, &idx, (uint32_t)model_len);
    append_le32(payload, &idx, model_offset_rel);

    // Entry: ExifIFDPointer (LONG)
    append_le16(payload, &idx, TAG_EXIF_IFD_POINTER);
    append_le16(payload, &idx, TYPE_LONG);
    append_le32(payload, &idx, 1);
    append_le32(payload, &idx, exif_ifd_offset_rel);

    // Entry: GPSInfoIFDPointer (LONG)
    append_le16(payload, &idx, TAG_GPS_IFD_POINTER);
    append_le16(payload, &idx, TYPE_LONG);
    append_le32(payload, &idx, 1);
    append_le32(payload, &idx, gps_ifd_offset_rel);

    // Next IFD offset (none)
    append_le32(payload, &idx, 0);

    // Write variable data: Make, Model strings
    // Position should now be tiff_start + make_offset_rel
    size_t expected_make_pos = tiff_start + make_offset_rel;
    if (idx < expected_make_pos) {
        size_t pad = expected_make_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    append_bytes(payload, &idx, make, make_len);

    size_t expected_model_pos = tiff_start + model_offset_rel;
    if (idx < expected_model_pos) {
        size_t pad = expected_model_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    append_bytes(payload, &idx, model, model_len);

    // Write ExifIFD (1 entry: ISO)
    size_t expected_exif_ifd_pos = tiff_start + exif_ifd_offset_rel;
    if (idx < expected_exif_ifd_pos) {
        size_t pad = expected_exif_ifd_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    append_le16(payload, &idx, 1);  // entry count
    // ISO entry (SHORT, count 1, value inline)
    append_le16(payload, &idx, TAG_ISO_SPEED_RATINGS);
    append_le16(payload, &idx, TYPE_SHORT);
    append_le32(payload, &idx, 1);
    // value/offset field: for SHORT count=1, put value in first 2 bytes (LE), then 0
    append_le16(payload, &idx, iso);
    append_le16(payload, &idx, 0);
    // next IFD offset
    append_le32(payload, &idx, 0);

    // Prepare GPS values
    char lat_ref = latitude >= 0.0 ? 'N' : 'S';
    char lon_ref = longitude >= 0.0 ? 'E' : 'W';
    double alat = fabs(latitude);
    double alon = fabs(longitude);
    uint32_t lat_deg = (uint32_t)floor(alat);
    double lat_min_f = (alat - lat_deg) * 60.0;
    uint32_t lat_min = (uint32_t)floor(lat_min_f);
    double lat_sec_f = (lat_min_f - lat_min) * 60.0;

    uint32_t lon_deg = (uint32_t)floor(alon);
    double lon_min_f = (alon - lon_deg) * 60.0;
    uint32_t lon_min = (uint32_t)floor(lon_min_f);
    double lon_sec_f = (lon_min_f - lon_min) * 60.0;

    // Write GPS IFD
    size_t expected_gps_ifd_pos = tiff_start + gps_ifd_offset_rel;
    if (idx < expected_gps_ifd_pos) {
        size_t pad = expected_gps_ifd_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    append_le16(payload, &idx, 5); // entry count

    // GPSVersionID (BYTE[4]) inline value {2,3,0,0}
    append_le16(payload, &idx, TAG_GPS_VERSION_ID);
    append_le16(payload, &idx, TYPE_BYTE);
    append_le32(payload, &idx, 4);
    // 4 bytes value inline
    append_bytes(payload, &idx, (uint8_t[4]){2,3,0,0}, 4);

    // LatitudeRef (ASCII[2]) inline 'N'/'S' + '\0'
    append_le16(payload, &idx, TAG_GPS_LATITUDE_REF);
    append_le16(payload, &idx, TYPE_ASCII);
    append_le32(payload, &idx, 2);
    append_bytes(payload, &idx, (uint8_t[4]){(uint8_t)lat_ref, 0, 0, 0}, 4);

    // Latitude (RATIONAL[3]) offset to data
    append_le16(payload, &idx, TAG_GPS_LATITUDE);
    append_le16(payload, &idx, TYPE_RATIONAL);
    append_le32(payload, &idx, 3);
    append_le32(payload, &idx, gps_lat_data_offset_rel);

    // LongitudeRef (ASCII[2]) inline 'E'/'W' + '\0'
    append_le16(payload, &idx, TAG_GPS_LONGITUDE_REF);
    append_le16(payload, &idx, TYPE_ASCII);
    append_le32(payload, &idx, 2);
    append_bytes(payload, &idx, (uint8_t[4]){(uint8_t)lon_ref, 0, 0, 0}, 4);

    // Longitude (RATIONAL[3]) offset to data
    append_le16(payload, &idx, TAG_GPS_LONGITUDE);
    append_le16(payload, &idx, TYPE_RATIONAL);
    append_le32(payload, &idx, 3);
    append_le32(payload, &idx, gps_lon_data_offset_rel);

    // next IFD offset for GPS
    append_le32(payload, &idx, 0);

    // Write GPS Latitude RATIONAL[3]: deg/min/sec
    size_t expected_lat_data_pos = tiff_start + gps_lat_data_offset_rel;
    if (idx < expected_lat_data_pos) {
        size_t pad = expected_lat_data_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    // degrees
    append_le32(payload, &idx, lat_deg);
    append_le32(payload, &idx, 1);
    // minutes
    append_le32(payload, &idx, lat_min);
    append_le32(payload, &idx, 1);
    // seconds (store with denominator 10000 for precision)
    append_le32(payload, &idx, (uint32_t)llround(lat_sec_f * 10000.0));
    append_le32(payload, &idx, 10000);

    // Write GPS Longitude RATIONAL[3]
    size_t expected_lon_data_pos = tiff_start + gps_lon_data_offset_rel;
    if (idx < expected_lon_data_pos) {
        size_t pad = expected_lon_data_pos - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    append_le32(payload, &idx, lon_deg);
    append_le32(payload, &idx, 1);
    append_le32(payload, &idx, lon_min);
    append_le32(payload, &idx, 1);
    append_le32(payload, &idx, (uint32_t)llround(lon_sec_f * 10000.0));
    append_le32(payload, &idx, 10000);

    return idx; // total payload size: ExifID + TIFF and all data
}

// Insert our APP1/Exif right after SOI, and skip any existing APP1 Exif segments while copying
static int write_exif_to_jpeg(const char *input_file,
                              const char *output_file,
                              const char *make,
                              const char *model,
                              double latitude,
                              double longitude,
                              uint16_t iso) {
    FILE *in = fopen(input_file, "rb");
    if (!in) { fprintf(stderr, "无法打开输入文件: %s\n", input_file); return -1; }
    FILE *out = fopen(output_file, "wb");
    if (!out) { fprintf(stderr, "无法创建输出文件: %s\n", output_file); fclose(in); return -1; }

    // Read and verify SOI
    uint8_t soi_bytes[2];
    if (fread(soi_bytes, 1, 2, in) != 2) { fprintf(stderr, "读取SOI失败\n"); fclose(in); fclose(out); return -1; }
    uint16_t soi = ((uint16_t)soi_bytes[0] << 8) | soi_bytes[1];
    if (soi != JPEG_SOI) { fprintf(stderr, "不是有效的JPEG文件\n"); fclose(in); fclose(out); return -1; }

    // Write SOI to output
    fwrite(soi_bytes, 1, 2, out);

    // Build our Exif payload in memory
    uint8_t payload[4096];
    size_t payload_size = build_exif_payload(payload, make, model, latitude, longitude, iso);

    // Write APP1 marker and length (length includes the 2 bytes of the length field itself)
    write_be16(out, JPEG_APP1);
    uint16_t app1_length = (uint16_t)(payload_size + 2);
    write_be16(out, app1_length);
    fwrite(payload, 1, payload_size, out);

    // Now copy input segments, skipping any existing APP1 Exif segments
    // Read segment by segment until SOS, then copy the rest raw
    for (;;) {
        int c1 = fgetc(in);
        if (c1 == EOF) break;
        if (c1 != 0xFF) {
            // Unexpected; write as raw and continue
            fputc(c1, out);
            continue;
        }
        int c2 = fgetc(in);
        if (c2 == EOF) break;
        // Write marker prefix 0xFF later depending on case
        uint16_t marker = ((uint16_t)0xFF << 8) | (uint8_t)c2;

        if (marker == JPEG_SOS) {
            // Write SOS marker
            write_be16(out, (uint16_t)marker);
            // Read length of SOS segment
            uint8_t len_bytes[2];
            if (fread(len_bytes, 1, 2, in) != 2) break;
            fwrite(len_bytes, 1, 2, out);
            uint16_t seg_len = ((uint16_t)len_bytes[0] << 8) | len_bytes[1];
            // seg_len includes its own 2 bytes; copy remaining seg_len-2 bytes of header
            if (seg_len >= 2) {
                uint16_t hdr_left = seg_len - 2;
                uint8_t *tmp = (uint8_t*)malloc(hdr_left);
                if (!tmp) { fclose(in); fclose(out); return -1; }
                if (fread(tmp, 1, hdr_left, in) != hdr_left) { free(tmp); break; }
                fwrite(tmp, 1, hdr_left, out);
                free(tmp);
            }
            // After SOS header, the rest until EOI is compressed data; copy raw until EOF
            uint8_t buf[4096];
            size_t n;
            while ((n = fread(buf, 1, sizeof buf, in)) > 0) {
                fwrite(buf, 1, n, out);
            }
            break;
        }

        // Read segment length (most markers have it; we assume standard JFIF structure)
        uint8_t len_bytes[2];
        if (fread(len_bytes, 1, 2, in) != 2) break;
        uint16_t seg_len = ((uint16_t)len_bytes[0] << 8) | len_bytes[1];
        uint16_t data_len = seg_len >= 2 ? seg_len - 2 : 0;

        // Read segment data
        uint8_t *seg_data = (uint8_t*)malloc(data_len);
        if (!seg_data) { fclose(in); fclose(out); return -1; }
        if (data_len > 0 && fread(seg_data, 1, data_len, in) != data_len) {
            free(seg_data);
            break;
        }

        int skip = 0;
        if (marker == JPEG_APP1 && data_len >= 6) {
            if (memcmp(seg_data, "Exif\0\0", 6) == 0) {
                // Skip existing Exif
                skip = 1;
            }
        }

        if (!skip) {
            write_be16(out, marker);
            write_be16(out, seg_len);
            if (data_len > 0) fwrite(seg_data, 1, data_len, out);
        }

        free(seg_data);
    }

    fclose(in);
    fclose(out);
    return 0;
}

int main(int argc, char* argv[]) {
    if (argc != 8) {
        printf("用法: %s <输入文件> <输出文件> <make> <model> <纬度> <经度> <ISO>\n", argv[0]);
        printf("示例: %s input.jpg output.jpg \"Canon\" \"EOS 5D Mark IV\" 39.9042 116.4074 100\n", argv[0]);
        return 1;
    }

    const char* input_file = argv[1];
    const char* output_file = argv[2];
    const char* make = argv[3];
    const char* model = argv[4];
    double latitude = atof(argv[5]);
    double longitude = atof(argv[6]);
    int iso = atoi(argv[7]);

    printf("正在向照片写入EXIF信息...\n");
    printf("Make: %s\n", make);
    printf("Model: %s\n", model);
    printf("GPS: %.6f, %.6f\n", latitude, longitude);
    printf("ISO: %d\n", iso);

    int rc = write_exif_to_jpeg(input_file, output_file, make, model, latitude, longitude, (uint16_t)iso);
    if (rc == 0) {
        printf("EXIF信息写入成功！输出文件: %s\n", output_file);
    } else {
        printf("EXIF信息写入失败！\n");
    }
    return rc;
} 